import pandas as pd
import os
import chardet


def load_data(file_path):
    try:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"文件不存在: {file_path}")

        with open(file_path, 'rb') as f:
            result = chardet.detect(f.read())

        data = pd.read_csv(
            file_path,
            encoding=result['encoding'],
            delimiter=',',
            usecols=['subject', 'relationship', 'capital'],  # 明确指定需要的列
            dtype=str
        )
        return data.dropna()  # 移除空值
    except Exception as e:
        log_path = os.path.join(os.path.dirname(__file__), '..', 'logs', 'error.log')
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        with open(log_path, 'a', encoding='utf-8') as log:
            log.write(f"加载数据失败: {str(e)}\n")
        raise


def process_data(data):
    try:
        # 统一处理字符串格式
        data['subject'] = data['subject'].str.strip()
        data['relationship'] = data['relationship'].str.strip()
        data['capital'] = data['capital'].str.strip()
        return data
    except Exception as e:
        log_path = os.path.join(os.path.dirname(__file__), '..', 'logs', 'error.log')
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        with open(log_path, 'a', encoding='utf-8') as log:
            log.write(f"处理数据失败: {str(e)}\n")
        raise


def save_data(data, output_file):
    try:
        os.makedirs(os.path.dirname(output_file), exist_ok=True)
        data.to_csv(output_file, index=False, encoding='utf-8')
    except Exception as e:
        log_path = os.path.join(os.path.dirname(__file__), '..', 'logs', 'error.log')
        os.makedirs(os.path.dirname(log_path), exist_ok=True)
        with open(log_path, 'a', encoding='utf-8') as log:
            log.write(f"保存数据失败: {str(e)}\n")
        raise